package cn.bby.query.server.infrastructure.config;

import cn.bby.query.server.domain.share.util.CrossDbUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.connect.runtime.WorkerConfig;
import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
import org.apache.kafka.connect.storage.Converter;
import org.apache.kafka.connect.storage.MemoryOffsetBackingStore;
import org.apache.kafka.connect.storage.OffsetUtils;

import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * @author wangbenyin
 * @version 1.0
 * @date 2025/07/02
 * @since 17
 */
@Slf4j
public class DbOffsetBackingStore extends MemoryOffsetBackingStore {
    private String topic;
    private final Map<String, Set<Map<String, Object>>> connectorPartitions;
    private final Converter keyConverter;

    public DbOffsetBackingStore(Converter keyConverter) {
        connectorPartitions = new HashMap<>();
        this.keyConverter = keyConverter;
    }

    @Override
    public Set<Map<String, Object>> connectorPartitions(String connectorName) {
        return connectorPartitions.getOrDefault(connectorName, Collections.emptySet());
    }

    @Override
    public void configure(WorkerConfig config) {
        topic = config.getString(DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG);
    }

    @Override
    public void start() {
        super.start();
        Map<byte[], byte[]> raw = CrossDbUtil.instance.queryOffset(topic);
        if (raw == null) {
            log.info("新建偏移量，topic:{}", topic);
            CrossDbUtil.instance.createOffset(topic);
            return;
        }
        for (Map.Entry<byte[], byte[]> mapEntry : raw.entrySet()) {
            ByteBuffer key = (mapEntry.getKey() != null) ? ByteBuffer.wrap(mapEntry.getKey()) : null;
            ByteBuffer value = (mapEntry.getValue() != null) ? ByteBuffer.wrap(mapEntry.getValue()) : null;
            data.put(key, value);
            OffsetUtils.processPartitionKey(mapEntry.getKey(), mapEntry.getValue(), keyConverter, connectorPartitions);
        }
    }

    @Override
    protected void save() {
        Map<byte[], byte[]> raw = new HashMap<>();
        for (Map.Entry<ByteBuffer, ByteBuffer> mapEntry : data.entrySet()) {
            byte[] key = (mapEntry.getKey() != null) ? mapEntry.getKey().array() : null;
            byte[] value = (mapEntry.getValue() != null) ? mapEntry.getValue().array() : null;
            raw.put(key, value);
            OffsetUtils.processPartitionKey(key, value, keyConverter, connectorPartitions);
        }
        CrossDbUtil.instance.updateOffset(topic, raw);
    }
}
